#!/usr/bin/env python3
"""
GPU 和 PyTorch 环境验证脚本
用于验证 PyTorch 1.13.1 是否正确安装并能识别 Quadro P6000 GPU
"""

import sys

def verify_pytorch():
    """验证 PyTorch 安装和 GPU 可用性"""
    
    print("=" * 70)
    print("🔍 PyTorch 环境验证")
    print("=" * 70)
    
    # 1. 导入 PyTorch
    try:
        import torch
        print("✅ PyTorch 导入成功")
    except ImportError as e:
        print(f"❌ PyTorch 导入失败: {e}")
        sys.exit(1)
    
    # 2. 检查 PyTorch 版本
    print(f"\n📦 PyTorch 版本: {torch.__version__}")
    expected_version = "1.13.1"
    if torch.__version__.startswith(expected_version):
        print(f"✅ 版本正确 (期望: {expected_version})")
    else:
        print(f"⚠️  版本不匹配 (期望: {expected_version}, 实际: {torch.__version__})")
    
    # 3. 检查 CUDA 是否可用
    print(f"\n🎮 CUDA 可用性: {torch.cuda.is_available()}")
    if not torch.cuda.is_available():
        print("❌ CUDA 不可用，无法使用 GPU 加速")
        print("   可能原因:")
        print("   1. PyTorch 未正确安装 CUDA 支持")
        print("   2. NVIDIA 驱动未安装")
        print("   3. GPU 不兼容")
        sys.exit(1)
    else:
        print("✅ CUDA 可用")
    
    # 4. 检查 CUDA 版本
    print(f"\n🔧 CUDA 版本: {torch.version.cuda}")
    
    # 5. 检查 cuDNN 版本
    if torch.backends.cudnn.is_available():
        print(f"✅ cuDNN 可用, 版本: {torch.backends.cudnn.version()}")
    else:
        print("⚠️  cuDNN 不可用")
    
    # 6. 检查 GPU 信息
    gpu_count = torch.cuda.device_count()
    print(f"\n🖥️  检测到 {gpu_count} 个 GPU:")
    
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        gpu_capability = torch.cuda.get_device_capability(i)
        gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
        
        print(f"\n  GPU {i}:")
        print(f"    名称: {gpu_name}")
        print(f"    计算能力: {gpu_capability[0]}.{gpu_capability[1]}")
        print(f"    显存: {gpu_memory:.2f} GB")
        
        # 验证计算能力
        if gpu_capability == (6, 1):
            print(f"    ✅ Quadro P6000 (Compute Capability 6.1) 已识别")
        else:
            print(f"    ⚠️  计算能力不是 6.1")
    
    # 7. 测试 GPU 计算
    print("\n🧪 测试 GPU 计算能力...")
    try:
        # 创建一个小张量并移到 GPU
        x = torch.randn(100, 100).cuda()
        y = torch.randn(100, 100).cuda()
        z = torch.matmul(x, y)
        
        print("✅ GPU 计算测试成功")
        print(f"   测试张量形状: {z.shape}")
        print(f"   测试张量设备: {z.device}")
        
    except Exception as e:
        print(f"❌ GPU 计算测试失败: {e}")
        sys.exit(1)
    
    # 8. 检查其他相关库
    print("\n📚 其他依赖库:")
    try:
        import torchvision
        print(f"  ✅ torchvision: {torchvision.__version__}")
    except ImportError:
        print("  ⚠️  torchvision 未安装")
    
    try:
        import torchaudio
        print(f"  ✅ torchaudio: {torchaudio.__version__}")
    except ImportError:
        print("  ⚠️  torchaudio 未安装")
    
    try:
        import numpy as np
        print(f"  ✅ numpy: {np.__version__}")
    except ImportError:
        print("  ⚠️  numpy 未安装")
    
    try:
        import pandas as pd
        print(f"  ✅ pandas: {pd.__version__}")
    except ImportError:
        print("  ⚠️  pandas 未安装")
    
    # 9. 总结
    print("\n" + "=" * 70)
    print("🎉 环境验证完成！")
    print("=" * 70)
    print("\n✅ 您的环境已正确配置，可以开始训练了！")
    print(f"✅ PyTorch {torch.__version__} 完全支持您的 Quadro P6000 GPU")
    print(f"✅ 可以使用 {gpu_count} 个 GPU 进行训练")
    print("\n💡 建议的训练命令:")
    print("   python src/train.py")
    print("\n")

if __name__ == "__main__":
    verify_pytorch()

